from transformers import (
    AutoTokenizer,
    AutoModelForCausalLM,
    BitsAndBytesConfig  # 用于量化，减少显存占用
)
import torch
import warnings
warnings.filterwarnings("ignore")  # 过滤无关警告（如弃用提示）

def load_gpt2_model(
    model_name: str = "gpt2",
    use_gpu: bool = True,
    load_in_8bit: bool = False,  # 8位量化（显存紧张时启用）
    cache_dir: str = "./huggingface_cache"  # 模型缓存路径（避免重复下载）
):
    """
    功能：加载 GPT-2 模型和分词器，支持 GPU 加速、量化优化
    参数：
        model_name: 模型名称（如 gpt2、gpt2-medium、gpt2-large）
        use_gpu: 是否使用 GPU（需提前安装 PyTorch GPU 版本）
        load_in_8bit: 是否启用 8 位量化（需安装 bitsandbytes 库）
        cache_dir: 模型缓存目录（默认本地 ./huggingface_cache）
    返回：
        tokenizer: 加载后的 GPT-2 分词器
        model: 加载后的 GPT-2 模型（已移动到指定设备）
    """
    # 1. 配置量化（显存不足时启用，如 8GB 显存加载 gpt2-large）
    quantization_config = None
    if load_in_8bit:
        quantization_config = BitsAndBytesConfig(
            load_in_8bit=True,
            bnb_8bit_use_double_quant=True,  # 双量化，进一步减少显存占用
            bnb_8bit_quant_type="nf4",       # 归一化浮点数，提升量化精度
            bnb_8bit_compute_dtype=torch.bfloat16  # 计算时用的精度
        )
        print(f"🔧 启用 8 位量化，模型显存占用将减少约 75%")

    # 2. 加载分词器（GPT-2 分词器需指定 padding_side 为右，避免生成时警告）
    try:
        tokenizer = AutoTokenizer.from_pretrained(
            model_name,
            cache_dir=cache_dir,
            padding_side="right",  # 右 padding（生成任务必需）
            use_fast=True  # 使用快速分词器（速度更快）
        )
        # GPT-2 原生无终止符，手动添加（避免生成无限文本）
        if tokenizer.pad_token is None:
            tokenizer.pad_token = tokenizer.eos_token  # 用 EOS 作为 PAD  token
        print(f"✅ 成功加载分词器：{model_name}")
    except Exception as e:
        raise RuntimeError(f"❌ 分词器加载失败：{str(e)}") from e

    # 3. 加载模型（自动判断 CPU/GPU）
    device = torch.device("cuda" if (use_gpu and torch.cuda.is_available()) else "cpu")
    try:
        model = AutoModelForCausalLM.from_pretrained(
            model_name,
            cache_dir=cache_dir,
            quantization_config=quantization_config,
            torch_dtype=torch.bfloat16 if device.type == "cuda" else torch.float32,  # GPU 用 bfloat16 提速
            device_map="auto"  # 自动分配设备（多 GPU 时自动拆分）
        )
        model.to(device)  # 确保模型移动到指定设备
        model.eval()  # 切换到评估模式（禁用 Dropout，生成更稳定）
        print(f"✅ 成功加载模型：{model_name}（设备：{device.type}，量化：{'8位' if load_in_8bit else '无'}）")
    except Exception as e:
        raise RuntimeError(f"❌ 模型加载失败：{str(e)}") from e

    return tokenizer, model, device

def generate_text(
    tokenizer,
    model,
    device,
    prompt: str,
    max_new_tokens: int = 100,  # 生成文本的最大长度（不含 prompt）
    temperature: float = 0.7,   # 随机性：0（确定）→1（随机）→2（极随机）
    top_p: float = 0.9,         # 核采样：只从概率前 top_p 的词中选（避免无意义词）
    repetition_penalty: float = 1.2,  # 重复惩罚：>1 减少重复内容（如“的的的”）
    do_sample: bool = True      # 是否用采样生成（False 为贪心搜索，速度快但生硬）
):
    """
    功能：用 GPT-2 生成文本，支持自定义生成策略
    参数：
        prompt: 输入提示词（如“人工智能的未来”）
        max_new_tokens: 生成文本的最大长度（建议不超过 512，GPT-2 上下文窗口为 1024）
        temperature: 控制生成随机性（推荐 0.6-0.9，故事生成用 0.8，事实问答用 0.3）
        top_p: 核采样阈值（推荐 0.8-0.95，与 temperature 二选一调优，避免同时过高）
        repetition_penalty: 重复惩罚（推荐 1.1-1.3，避免生成重复句子）
        do_sample: 是否采样（True 生成更自然，False 速度更快）
    返回：
        generated_text: 最终生成的完整文本（prompt + 生成内容）
    """
    # 1. 对输入提示词分词（返回 PyTorch 张量，并移动到模型设备）
    inputs = tokenizer(
        prompt,
        return_tensors="pt",  # 返回 PyTorch 张量
        padding=True,
        truncation=True,      # 截断过长的 prompt（避免超出 GPT-2 1024 上下文窗口）
        max_length=1024 - max_new_tokens  # 预留生成空间
    ).to(device)

    # 2. 生成文本（禁用梯度计算，提升速度并减少显存占用）
    with torch.no_grad():
        outputs = model.generate(
            **inputs,
            max_new_tokens=max_new_tokens,
            temperature=temperature,
            top_p=top_p,
            repetition_penalty=repetition_penalty,
            do_sample=do_sample,
            eos_token_id=tokenizer.eos_token_id,  # 遇到 EOS 停止生成
            pad_token_id=tokenizer.pad_token_id,  # 处理 padding 的 token
            no_repeat_ngram_size=3  # 避免 3 个词以上的重复（如“今天天气好，今天天气好”）
        )

    # 3. 解码张量为文本（跳过特殊 token，如 PAD/EOS）
    generated_text = tokenizer.decode(
        outputs[0],
        skip_special_tokens=True,  # 跳过 <|endoftext|> 等特殊 token
        clean_up_tokenization_spaces=True  # 清理多余空格
    )

    return generated_text

# ------------------- 示例调用 -------------------
if __name__ == "__main__":
    # 1. 加载模型（根据显存调整：8GB 显存建议用 gpt2 + load_in_8bit=True；16GB 可加载 gpt2-large）
    tokenizer, model, device = load_gpt2_model(
        model_name="gpt2",  # 可选：gpt2-medium（3.5B参数）、gpt2-large（7.7B参数）
        use_gpu=True,
        load_in_8bit=False,  # 若显存不足，改为 True（需安装 pip install bitsandbytes）
        cache_dir="./huggingface_cache"
    )

    # 2. 定义输入提示词（根据场景自定义，如故事开头、问题、创意文案）
    user_prompt = "请写一段关于‘人工智能如何帮助保护濒危动物’的短文，开头为：‘在非洲的塞伦盖蒂草原上，一头年迈的长颈鹿正艰难地寻找水源。’"

    # 3. 生成文本
    print(f"\n📝 输入提示词：{user_prompt}")
    print(f"\n🔄 正在生成文本...")
    result = generate_text(
        tokenizer=tokenizer,
        model=model,
        device=device,
        prompt=user_prompt,
        max_new_tokens=150,  # 生成 150 个词
        temperature=0.75,    # 中等随机性，兼顾自然和逻辑
        top_p=0.92,
        repetition_penalty=1.2
    )

    # 4. 打印结果（格式化输出，区分 prompt 和生成内容）
    print(f"\n✅ 生成完成！")
    print("="*50)
    # 提取生成的内容（排除原 prompt）
    prompt_length = len(tokenizer.encode(user_prompt, add_special_tokens=False))
    generated_content = tokenizer.decode(
        tokenizer.encode(result, add_special_tokens=False)[prompt_length:],
        skip_special_tokens=True
    )
    print(f"原始提示：{user_prompt}")
    print(f"生成内容：{generated_content}")
    print("="*50)